import numpy as np
import torch
import datetime
from typing import Tuple
from rich import print


class EgoMotionComp:
    imageShape = (800, 1280)
    z = 1758.38

    # Parameters
    a = 0.1
    b = 0.455

    def __init__(self, scale: int = 1, device = None) -> None:
        """
        自身运动补偿。
        @params:
            scale: int 缩放比例
            device: torch.device 计算设备
        """
        self.scale = scale
        self.device = device
        self.reset()

    def reset(self) -> None:
        """
        重置该单元。
        """
        self.k = torch.zeros(3, dtype = torch.float, device = self.device)
        self.omegaM = 0.0
        self.processTimes = [0.0] * 2
        self.lastR = None
        self.lastTime = 0.0
        self.duration1 = 0.0
        self.duration2 = 0.0

    def time(self) -> float:
        """
        获取当前时间
        @return:
            double 当前的时间戳（单位为秒）
        """
        res = datetime.datetime.now().timestamp()
        return res
    
    def updateIMUTime(self, currentTime: float) -> None:
        """
        更新IMU时间。
        @params:
            currentTime: float 当前时间
        """
        if self.lastTime == 0.0:
            self.duration1 = 0.0
        self.duration2 = self.duration1
        self.duration1 = currentTime - self.lastTime
        self.lastTime = currentTime

    def processImu(self, data: torch.Tensor) -> None:
        """
        处理IMU数据
        @params:
            imu: torch.Tensor IMU数据
        """

        if not data.shape[0]:
            return
        #由传入IMU信息计算角速度大小
        deltaTime = self.duration1 / data.shape[0]
        omegaZ = data[:, 1] # [n]
        omegaY = data[:, 2] # [n]
        omegaX = data[:, 3] # [n]
        xDegrees = omegaX * deltaTime
        yDegrees = omegaY * deltaTime
        zDegrees = omegaZ * deltaTime
        Rx = torch.zeros(data.shape[0], 3, 3, dtype = torch.float, device = self.device) # [n, 3, 3]
        Rx[:, 0, 0] = 1.0
        Rx[:, 1, 1] = Rx[:, 2, 2] = torch.cos(xDegrees)
        Rx[:, 1, 2] = -torch.sin(xDegrees)
        Rx[:, 2, 1] = torch.sin(xDegrees)
        Ry = torch.zeros(data.shape[0], 3, 3, dtype = torch.float, device = self.device) # [n, 3, 3]
        Ry[:, 0, 0] = Ry[:, 2, 2] = torch.cos(yDegrees)
        Ry[:, 0, 2] = torch.sin(yDegrees)
        Ry[:, 1, 1] = 1.0
        Ry[:, 2, 0] = -torch.sin(yDegrees)
        Rz = torch.zeros(data.shape[0], 3, 3, dtype = torch.float, device = self.device) # [n, 3, 3]
        Rz[:, 0, 0] = Rz[:, 1, 1] = torch.cos(zDegrees)
        Rz[:, 0, 1] = -torch.sin(zDegrees)
        Rz[:, 1, 0] = torch.sin(zDegrees)
        Rz[:, 2, 2] = 1.0
        RArr = Rz @ Ry @ Rx # [n, 3, 3]
        R = torch.eye(3, dtype = torch.float, device = self.device)
        for tempR in range(RArr.shape[0]):
            R = RArr[tempR] @ R
        if self.lastR is None:
            self.lastR = R.clone()
        else:
            lastR = R.clone()
            R = R @ self.lastR
            self.lastR = lastR
        k, theta = self.getTheta(R)
        duration = self.duration1 + self.duration2
        self.omegaM = theta / duration if duration else 0.0
        self.k = k

    def getTheta(self, R: torch.Tensor) -> Tuple[torch.Tensor, float]:
        """        
        获取转轴K和角度值theta
        @params:
            R: torch.Tensor 旋转矩阵
        @return:
            k: torch.Tensor 转轴
            theta: double 转角theta
        """
        k = torch.tensor([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
        k_norm = torch.norm(k) + 1e-9
        k = k / k_norm
        theta = torch.arcsin(k_norm / 2)
        
        return k, theta

    def getR(self, k: torch.Tensor, theta: torch.Tensor) -> torch.Tensor:
        """
        重建旋转矩阵R
        @params:
            k: torch.Tensor 转轴
            theta: torch.Tensor 旋转角度
        @return:
            torch.Tensor 旋转矩阵
        """
        K = torch.tensor([[0.0, -k[2], k[1]], [k[2], 0.0, -k[0]], [-k[1], k[0], 0.0]], dtype = torch.float, device = self.device)
        I = torch.eye(3, dtype = torch.float, device = self.device)
        R = torch.einsum("mn,b->bmn", I, torch.ones_like(theta)) + torch.einsum("mn,b->bmn", K, torch.sin(theta)) + torch.einsum("mn,b->bmn", K @ K, 1.0 - torch.cos(theta))
        return R

    def averageTime(self, xy: torch.Tensor, idx: torch.Tensor, t: torch.Tensor, xSize: int, ySize: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        统计每个坐标的平均时间，返回统计后的量。
        @params:
            xy: torch.Tensor 空间坐标，使用x * height + y来表示，形状为[p]，即经过排除重复之后的不同元素长度
            idx: torch.Tensor 原来每个事件的x和y对应新的空间坐标的下标，形状为[n]
            t: torch.Tensor 原来每个事件的时间戳，形状为[n]
        @return:
            avrT: torch.Tensor 各个事件对应的平均时间，形状为[n]
            globalAvrT: torch.Tensor 总平均时间
        """
        pointCount = xy.shape[0]
        eventCount = idx.shape[0]
        mask = torch.zeros(pointCount, eventCount).to(t)
        mask = (mask.T + torch.arange(pointCount).to(t)).T - idx
        mask = (mask == 0).to(t)
        timeMap = torch.sum(mask * t, dim = 1)
        countMap = torch.sum(mask, dim = 1)
        averageT = torch.where(countMap != 0.0, timeMap / countMap, timeMap)
        res = (mask.T * averageT).T
        res = torch.sum(res, dim = 0)
        globalRes = torch.sum(res)
        if pointCount:
            globalRes /= pointCount
        return res, globalRes

    def execute(self, data: torch.Tensor) -> torch.Tensor:
        """
        根据所设定的参数，执行自身运动补偿
        @params:
            data: torch.Tensor 未进行运动补偿的事件列表
        @return:
            torch.Tensor 运动补偿后的事件列表
        """
        processStart = self.time()
        res = torch.tensor([], dtype = torch.long, device = self.device)
        if not data.shape[0]:
            return res

        xSize = self.imageShape[1] // self.scale
        ySize = self.imageShape[0] // self.scale
        radius = self.z
        calcRatio = lambda x, y, z, r: r / torch.sqrt((x ** 2) + (y ** 2) + (z ** 2))
        xCenter = (xSize - 1.0) / 2.0
        yCenter = (ySize - 1.0) / 2.0

        if data.shape[0]:
            sec = data[:, 0].to(torch.float)
            sec -= torch.min(sec)
            nsec = data[:, 1].to(torch.float)
            tStart = sec + nsec * 1e-9
            minTimeStamp = torch.min(tStart)
            maxTimeStamp = torch.max(tStart)
            newEvent = torch.ones((data.shape[0], 3), dtype = torch.float, device = self.device)
            newEvent[:, 1] = data[:, 3].to(torch.float) - yCenter # y加个偏移量, 大小为y的长度（800）-1后/2
            newEvent[:, 0] = (xSize - 1.0 - data[:, 4].to(torch.float)) - xCenter # x加个偏移量, 大小为x的长度（1280）-1后/2
            ratio = calcRatio(newEvent[:, 0], newEvent[:, 1], self.z, radius)
            newEvent[:, 0] = newEvent[:, 0] * ratio
            newEvent[:, 1] = newEvent[:, 1] * ratio
            newEvent[:, 2] = self.z * ratio
            
            deltaTime = tStart - minTimeStamp
            RInv = self.getR(self.k, -self.omegaM * deltaTime)
            newEvent = torch.einsum("bmn,bn->bm", RInv, newEvent)

            ratio = newEvent[:, 2] / radius
            ratio = 1.0 * (ratio == 0.0).to(torch.float) + ratio
            newEvent[:, 0] = newEvent[:, 0] / ratio
            newEvent[:, 1] = newEvent[:, 1] / ratio
            newEvent[:, 2] = newEvent[:, 2] / ratio

            yxt = torch.zeros_like(newEvent)
            yxt[:, 0] = torch.round(newEvent[:, 1] + yCenter) # y
            yxt[:, 1] = torch.round(newEvent[:, 0] + xCenter) # x
            yxt[:, 2] = deltaTime.clone()
            frameFilter = (yxt[:, 0] >= 0) & (yxt[:, 0] < ySize) & (yxt[:, 1] >= 0) & (yxt[:, 1] < xSize)
            data = data[frameFilter]
            yxt = yxt[frameFilter]

            if not data.shape[0]:
                return res

            newIdx = yxt[:, 0] * int(xSize) + yxt[:, 1]

            # 检查点1
            checkPoint1 = self.time()
            self.processTimes[0] = checkPoint1 - processStart

            events, eventIdx = torch.unique(newIdx, return_inverse = True)
            avrTime = torch.zeros(eventIdx.shape[0], dtype = torch.float, device = self.device)
            avr = torch.tensor(0, dtype = torch.float, device = self.device)
            avrTime, avr = self.averageTime(events, eventIdx, yxt[:, 2], int(xSize), int(ySize))
            yxt[:, 2] = torch.abs((avrTime - avr) / (maxTimeStamp - minTimeStamp)) if maxTimeStamp != minTimeStamp else 0.0
            # 得到yxt: 变换后的事件序列，形状为[n, 3]，第一维是x，第二维是y，第三维是算出来的归一化时间值
            # yxt是中间变量，需要最后使用yxt来排除掉原本的事件不符合要求的那些
            threshold = self.a * self.omegaM + self.b

            res = data[yxt[:, 2] >= threshold]
            
            # 检查点2
            checkPoint2 = self.time()
            self.processTimes[1] = checkPoint2 - checkPoint1

        return res